#
# mydbms/pysql_compiler/planner.py
#

from .ast import (
    ASTNode, SelectStatementNode, BinaryOperationNode, IdentifierNode,
    NumberNode, StringNode, StarNode, InsertStatementNode, UpdateStatementNode,
    DeleteStatementNode, AssignmentNode, CreateTableStatementNode,
    ColumnDefinitionNode, DataTypeNode, CreateIndexStatementNode, OrderByNode,
    FunctionCallNode, AliasNode, TruncateTableStatementNode, InPredicateNode, JoinNode,
    # <<<--- 导入约束节点 (主键约束等) ---<<<
    PrimaryKeyConstraintNode, AutoIncrementConstraintNode, NotNullConstraintNode,
    UniqueConstraintNode, CheckConstraintNode, DropTableStatementNode,
    AlterTableStatementNode, AddColumnNode, AlterColumnNode, DropColumnNode, AddPrimaryKeyNode,
    UnaryOperationNode, LikePredicateNode, IsNullPredicateNode, DropIndexStatementNode,
    ShowTablesStatementNode,
    ExplainStatementNode, BooleanLiteralNode
)
from .catalog import Catalog


class Planner:
    """
    遍历AST并生成查询计划。
    """

    def __init__(self, catalog: Catalog):
        self.catalog = catalog

    def create_plan(self, ast_node: ASTNode) -> dict:
        if isinstance(ast_node, SelectStatementNode):
            return self._plan_select_statement(ast_node)
        elif isinstance(ast_node, InsertStatementNode):
            plans = self._plan_insert_statement(ast_node)
            if len(plans) == 1:
                return plans[0]
            return plans
        elif isinstance(ast_node, UpdateStatementNode):
            return self._plan_update_statement(ast_node)
        elif isinstance(ast_node, DeleteStatementNode):
            return self._plan_delete_statement(ast_node)
        elif isinstance(ast_node, CreateTableStatementNode):
            return self._plan_create_table_statement(ast_node)
        elif isinstance(ast_node, CreateIndexStatementNode):
            return self._plan_create_index_statement(ast_node)
        elif isinstance(ast_node, TruncateTableStatementNode):
            return self._plan_truncate_table_statement(ast_node)
        elif isinstance(ast_node, DropTableStatementNode):
            return self._plan_drop_table_statement(ast_node)
        elif isinstance(ast_node, DropIndexStatementNode):
            return self._plan_drop_index_statement(ast_node)
        #处理ALTER语句
        elif isinstance(ast_node, AlterTableStatementNode):
            return self._plan_alter_table_statement(ast_node)
        elif isinstance(ast_node, ShowTablesStatementNode):
            return self._plan_show_tables_statement(ast_node)
        elif isinstance(ast_node, ExplainStatementNode):
            inner_plan = self.create_plan(ast_node.statement)
            return {
            'operation': 'EXPLAIN',
            'plan': inner_plan
            }
        else:
            raise ValueError(f"Cannot create plan for AST node type: {type(ast_node).__name__}")

    def _plan_truncate_table_statement(self, node: TruncateTableStatementNode) -> dict:
        return {
            'operation': 'TRUNCATE_TABLE',
            'table_name': node.table_name.value
        }

    def _plan_select_statement(self, node: SelectStatementNode) -> dict:
        # <<<--- 核心修改：处理 from_clause ---<<<
        plan = self._plan_from_clause(node.from_clause)
        # >>>--- 修改结束 ---<<<

        if node.where_clause:
            plan = {
                'operation': 'FILTER',
                'predicate': self._plan_expression(node.where_clause),
                'source': plan
            }

        # <<<--- 修正：确保聚合函数能被正确识别 ---<<<
        # 检查列中是否有聚合函数调用
        is_aggregate = any(
            isinstance(c, FunctionCallNode) or (isinstance(c, AliasNode) and isinstance(c.expression, FunctionCallNode))
            for c in node.columns
        )
        if node.group_by_clause or is_aggregate:
            output_schema = []
            for col in node.columns:
                if isinstance(col, AliasNode):
                    output_schema.append(col.alias.value)
                elif isinstance(col, IdentifierNode):
                    output_schema.append(col.value)
                # <<<--- 修正：为函数调用（无别名）生成默认名 ---<<<
                elif isinstance(col, FunctionCallNode):
                    arg_str = ''
                    if col.args:
                        # 检查参数是否是 * (StarNode)
                        if isinstance(col.args[0], StarNode):
                            arg_str = '*'
                        else:
                            arg_str = col.args[0].value
                    output_schema.append(f"{col.func_name}({arg_str})")
                # >>>--- 修正结束 ---<<<

            plan = {
                'operation': 'AGGREGATE',
                'group_by': [self._plan_expression(col) for col in
                             node.group_by_clause] if node.group_by_clause else [],
                'aggregates': [self._plan_expression(col) for col in node.columns],
                'output_schema': output_schema,
                'source': plan
            }

            if node.having_clause:
                plan = {
                    'operation': 'FILTER',
                    'is_having': True,
                    'predicate': self._plan_expression(node.having_clause),
                    'source': plan
                }

        plan = {
            'operation': 'PROJECT',
            'columns': [self._plan_expression(col) for col in node.columns],
            'source': plan
        }

        if node.distinct:
            plan = {
                'operation': 'DISTINCT',
                'source': plan
            }

        if node.order_by_clause:
            plan = {
                'operation': 'SORT',
                'sort_key': node.order_by_clause.column.value,
                'order': node.order_by_clause.order,
                'source': plan
            }

        return plan

    def _plan_from_clause(self, node):
        """为 FROM 子句（可能包含 JOIN）生成计划"""
        if isinstance(node, JoinNode):
            left_plan = self._plan_from_clause(node.left)
            right_plan = self._plan_from_clause(node.right)
            return {
                'operation': 'JOIN',
                'join_type': node.join_type,
                'left_source': left_plan,
                'right_source': right_plan,
                'condition': self._plan_expression(node.condition)
            }
        elif isinstance(node, IdentifierNode):
            table_name = node.value
            table_schema = self.catalog.get_table_schema(table_name)
            if not table_schema:
                raise ValueError(f"Table '{table_name}' not found during planning.")
            return {'operation': 'SCAN', 'table_name': table_name}
        else:
            raise ValueError(f"Unsupported FROM clause node type: {type(node)}")

    def _plan_insert_statement(self, node: InsertStatementNode) -> list:
        plans = []
        for value_list in node.values:
            plan = {
                'operation': 'INSERT',
                'table_name': node.table.value,
                'columns': [col.value for col in node.columns] if node.columns else [],
                'values': [self._plan_expression(val) for val in value_list]
            }
            plans.append(plan)
        return plans

    def _plan_update_statement(self, node: UpdateStatementNode) -> dict:
        plan = {
            'operation': 'UPDATE',
            'table_name': node.table.value,
            'assignments': [self._plan_expression(ass) for ass in node.assignments]
        }
        if node.from_table:
            plan['from_table'] = self._plan_expression(node.from_table)
        if node.where_clause:
            plan['predicate'] = self._plan_expression(node.where_clause)
        return plan

    def _plan_delete_statement(self, node: DeleteStatementNode) -> dict:
        plan = {
            'operation': 'DELETE',
            'table_name': node.table.value
        }
        if node.where_clause:
            plan['predicate'] = self._plan_expression(node.where_clause)
        return plan

    def _plan_create_table_statement(self, node: CreateTableStatementNode) -> dict:
        columns = []
        constraints = []
        # 遍历AST节点，将列定义和表约束分开放入不同列表
        for definition in node.columns:
            if isinstance(definition, ColumnDefinitionNode):
                columns.append(self._plan_expression(definition))
            elif isinstance(definition, UniqueConstraintNode):
                constraints.append(self._plan_expression(definition))

        return {
            'operation': 'CREATE_TABLE',
            'table_name': node.table_name.value,
            'columns': columns,
            'constraints': constraints  # 在计划中新增一个约束字段
        }

    def _plan_create_index_statement(self, node: CreateIndexStatementNode) -> dict:
        return {
            'operation': 'CREATE_INDEX',
            'index_name': node.index_name.value,
            'table_name': node.table_name.value,
            'column_name': node.column_name.value
        }

    def _plan_expression(self, node: ASTNode):
        if isinstance(node, InPredicateNode):
            # 修改以处理值列表和子查询
            if isinstance(node.subquery, list):
                return {
                    'op': 'IN',
                    'identifier': self._plan_expression(node.identifier),
                    'values': [self._plan_expression(val) for val in node.subquery]
                }
            else: # is a subquery
                return {
                    'op': 'IN',
                    'identifier': self._plan_expression(node.identifier),
                    'subquery_plan': self.create_plan(node.subquery)
                }
        # <<<--- 新增: 为新的 AST 节点生成计划 ---<<<
        elif isinstance(node, UnaryOperationNode):
            return {'op': node.op, 'operand': self._plan_expression(node.operand)}
        elif isinstance(node, LikePredicateNode):
            return {'op': 'LIKE', 'left': self._plan_expression(node.left), 'pattern': self._plan_expression(node.pattern)}
        elif isinstance(node, IsNullPredicateNode):
            return {'op': 'IS_NULL', 'left': self._plan_expression(node.left), 'is_not': node.is_not}
        # >>>--- 结束 ---<<<
        if isinstance(node, BinaryOperationNode):
            return {'type': 'binary_op','op': node.op, 'left': self._plan_expression(node.left), 'right': self._plan_expression(node.right)}
        elif isinstance(node, AssignmentNode):
            return {'column': self._plan_expression(node.column), 'value': self._plan_expression(node.value)}
        elif isinstance(node, ColumnDefinitionNode):
            # <<<--- 修改：在计划中包含约束信息 (主键约束) ---<<<
            return {'column_name': self._plan_expression(node.column_name),
                    'data_type': self._plan_expression(node.data_type),
                    'constraints': [self._plan_expression(c) for c in node.constraints]}
            # >>>--- 修改结束 ---<<<
        elif isinstance(node, DataTypeNode):
            # 原始代码: return {'type': 'data_type', 'name': node.value}

            # --- 替换为以下代码 ---
            full_type_string = node.value
            params = []
            if '(' in full_type_string and full_type_string.endswith(')'):
                type_name, param_part = full_type_string.split('(', 1)
                param_values = param_part[:-1].split(',')
                # 简单处理，转换为整数，更复杂的可以处理其他类型
                params = [int(p.strip()) for p in param_values]
            else:
                type_name = full_type_string

            return {'type': 'data_type', 'name': type_name, 'params': params}
            # --- 替换结束 ---

        elif isinstance(node, IdentifierNode):
            return {'type': 'identifier', 'name': node.value}
        elif isinstance(node, (NumberNode, StringNode, BooleanLiteralNode)):  # 在元组中加入 BooleanLiteralNode
            return {'type': 'literal', 'value': node.value}
        elif isinstance(node, StarNode):
            return {'type': 'star'}
        elif isinstance(node, OrderByNode):
            return {'type': 'order_by', 'column': node.column.value, 'order': node.order}
        elif isinstance(node, FunctionCallNode):
            # <<<--- 新增：处理 COUNT(*) 的特殊情况 ---<<<
            args_plan = []
            if node.args and isinstance(node.args[0], StarNode):
                args_plan.append({'type': 'star'})
            else:
                args_plan = [self._plan_expression(arg) for arg in node.args]
            return {'type': 'function_call', 'func_name': node.func_name,
                    'args': args_plan}
        elif isinstance(node, AliasNode):
            return {'type': 'alias', 'expression': self._plan_expression(node.expression),
                    'alias': self._plan_expression(node.alias)}
            # <<<--- 处理约束节点 (主键约束等) ---<<<
        elif isinstance(node, PrimaryKeyConstraintNode):
            return {'type': 'constraint', 'name': 'PRIMARY KEY'}
        elif isinstance(node, AutoIncrementConstraintNode):
            return {'type': 'constraint', 'name': 'AUTO_INCREMENT'}
        elif isinstance(node, NotNullConstraintNode):
            return {'type': 'constraint', 'name': 'NOT NULL'}
        elif isinstance(node, UniqueConstraintNode):
            # 如果 node.columns 存在，说明是多列约束
            if node.columns:
                return {'type': 'constraint', 'name': 'UNIQUE', 'columns': [c.value for c in node.columns]}
            # 否则，是单列约束
            else:
                return {'type': 'constraint', 'name': 'UNIQUE'}
        elif isinstance(node, CheckConstraintNode):
            return {'type': 'constraint', 'name': 'CHECK', 'expression': self._plan_expression(node.expression)}
        # >>>--- 新增结束 ---<<<
        else:
            raise ValueError(f"Unknown expression node type: {type(node).__name__}")

    def _plan_drop_table_statement(self, node: DropTableStatementNode) -> dict:
        return {
            'operation': 'DROP_TABLE',
            'table_name': node.table_name.value
        }

    # <<<--- 为 ALTER TABLE 生成计划 ---<<<
    def _plan_alter_table_statement(self, node: AlterTableStatementNode) -> dict:
        plan = {
            'operation': 'ALTER_TABLE',
            'table_name': node.table_name.value
        }
        action_node = node.action

        # 创建一个字典来封装 action 的所有信息
        action_plan = {}
        if isinstance(action_node, AddColumnNode):
            action_plan['type'] = 'ADD_COLUMN'  # 键名为 'type'
            action_plan['column_def'] = self._plan_expression(action_node.column_definition)
        elif isinstance(action_node, AlterColumnNode):
            action_plan['type'] = 'ALTER_COLUMN'
            action_plan['column_name'] = action_node.column_name.value
            action_plan['new_data_type'] = self._plan_expression(action_node.new_data_type)
        elif isinstance(action_node, DropColumnNode):
            action_plan['type'] = 'DROP_COLUMN'
            action_plan['column_name'] = action_node.column_name.value
        elif isinstance(action_node, AddPrimaryKeyNode):
            action_plan['type'] = 'ADD_PRIMARY_KEY'
            action_plan['column_name'] = action_node.column_name.value

        # 将封装好的 action 字典赋值给 plan 的 'action' 键
        plan['action'] = action_plan
        return plan

    def _plan_drop_index_statement(self, node: DropIndexStatementNode) -> dict:
        """为 DROP INDEX 生成计划"""
        return {
            'operation': 'DROP_INDEX',
            'index_name': node.index_name.value
        }

    def _plan_show_tables_statement(self, node: ShowTablesStatementNode) -> dict:
        """为 SHOW TABLES 生成计划"""
        return {
            'operation': 'SHOW_TABLES'
        }